{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center>\n",
    "    <p style=\"text-align:center\">\n",
    "        <img alt=\"phoenix logo\" src=\"https://raw.githubusercontent.com/Arize-ai/phoenix-assets/9e6101d95936f4bd4d390efc9ce646dc6937fb2d/images/socal/github-large-banner-phoenix.jpg\" width=\"1000\"/>\n",
    "        <br>\n",
    "        <br>\n",
    "        <a href=\"https://arize.com/docs/phoenix/\">Docs</a>\n",
    "        |\n",
    "        <a href=\"https://github.com/Arize-ai/phoenix\">GitHub</a>\n",
    "        |\n",
    "        <a href=\"https://arize-ai.slack.com/join/shared_invite/zt-2w57bhem8-hq24MB6u7yE_ZF_ilOYSBw#/shared-invite/email\">Community</a>\n",
    "    </p>\n",
    "</center>\n",
    "<h1 align=\"center\">Tracing CrewAI with Arize Phoenix - Orchestrator-Workers Workflow</h1>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q arize-phoenix opentelemetry-sdk opentelemetry-exporter-otlp crewai crewai_tools openinference-instrumentation-crewai"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5-gPdVmIndw9"
   },
   "source": [
    "# Set up Keys and Dependencies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: For this colab you'll need:\n",
    "\n",
    "*   OpenAI API key (https://openai.com/)\n",
    "*   Serper API key (https://serper.dev/)\n",
    "*   Phoenix API key (https://app.phoenix.arize.com/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from getpass import getpass\n",
    "\n",
    "# Prompt the user for their API keys if they haven't been set\n",
    "openai_key = os.getenv(\"OPENAI_API_KEY\", \"OPENAI_API_KEY\")\n",
    "serper_key = os.getenv(\"SERPER_API_KEY\", \"SERPER_API_KEY\")\n",
    "\n",
    "if openai_key == \"OPENAI_API_KEY\":\n",
    "    openai_key = getpass(\"Please enter your OPENAI_API_KEY: \")\n",
    "\n",
    "if serper_key == \"SERPER_API_KEY\":\n",
    "    serper_key = getpass(\"Please enter your SERPER_API_KEY: \")\n",
    "\n",
    "# Set the environment variables with the provided keys\n",
    "os.environ[\"OPENAI_API_KEY\"] = openai_key\n",
    "os.environ[\"SERPER_API_KEY\"] = serper_key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if \"PHOENIX_API_KEY\" not in os.environ:\n",
    "    os.environ[\"PHOENIX_API_KEY\"] = getpass(\"🔑 Enter your Phoenix API key: \")\n",
    "\n",
    "if \"PHOENIX_COLLECTOR_ENDPOINT\" not in os.environ:\n",
    "    os.environ[\"PHOENIX_COLLECTOR_ENDPOINT\"] = getpass(\"🔑 Enter your Phoenix Collector Endpoint\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "r9X87mdGnpbc"
   },
   "source": [
    "## Configure Tracing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from phoenix.otel import register\n",
    "\n",
    "tracer_provider = register(project_name=\"crewai-agents\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vYT-EU56ni94"
   },
   "source": [
    "# Instrument CrewAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openinference.instrumentation.crewai import CrewAIInstrumentor\n",
    "\n",
    "CrewAIInstrumentor().instrument(skip_dep_check=True, tracer_provider=tracer_provider)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define your Agents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from crewai import Agent, Crew, Task\n",
    "\n",
    "# Define worker agents\n",
    "trend_researcher = Agent(\n",
    "    role=\"AI Trend Researcher\",\n",
    "    goal=\"Analyze current advancements in AI\",\n",
    "    backstory=\"Expert in tracking and analyzing new trends in artificial intelligence.\",\n",
    "    verbose=True,\n",
    ")\n",
    "\n",
    "policy_analyst = Agent(\n",
    "    role=\"AI Policy Analyst\",\n",
    "    goal=\"Examine the implications of AI regulations and governance\",\n",
    "    backstory=\"Tracks AI policy developments across governments and organizations.\",\n",
    "    verbose=True,\n",
    ")\n",
    "\n",
    "risk_specialist = Agent(\n",
    "    role=\"AI Risk Specialist\",\n",
    "    goal=\"Identify potential risks in frontier AI development\",\n",
    "    backstory=\"Focuses on safety, alignment, and misuse risks related to advanced AI.\",\n",
    "    verbose=True,\n",
    ")\n",
    "\n",
    "synthesizer = Agent(\n",
    "    role=\"Synthesis Writer\",\n",
    "    goal=\"Summarize all findings into a final cohesive report\",\n",
    "    backstory=\"Expert at compiling research insights into executive-level narratives.\",\n",
    "    verbose=True,\n",
    ")\n",
    "\n",
    "orchestrator = Agent(\n",
    "    role=\"Orchestrator\",\n",
    "    goal=(\n",
    "        \"Your job is to delegate research and writing tasks to the correct coworker using the 'Delegate work to coworker' tool.\\n\"\n",
    "        \"For each task you assign, you MUST call the tool with the following JSON input:\\n\\n\"\n",
    "        \"{\\n\"\n",
    "        '  \"task\": \"Short summary of the task to do (plain string)\",\\n'\n",
    "        '  \"context\": \"Why this task is important or part of the report (plain string)\",\\n'\n",
    "        '  \"coworker\": \"One of: AI Trend Researcher, AI Policy Analyst, AI Risk Specialist, Synthesis Writer\"\\n'\n",
    "        \"}\\n\\n\"\n",
    "        \"IMPORTANT:\\n\"\n",
    "        \"- Do NOT format 'task' or 'context' as dictionaries.\\n\"\n",
    "        \"- Do NOT include types or nested descriptions.\\n\"\n",
    "        \"- Only use plain strings for both.\\n\"\n",
    "        \"- Call the tool multiple times, one per coworker.\"\n",
    "    ),\n",
    "    backstory=\"You are responsible for assigning each part of an AI report to the right specialist.\",\n",
    "    verbose=True,\n",
    "    allow_delegation=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define your Tasks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the initial task only for the orchestrator\n",
    "initial_task = Task(\n",
    "    description=\"Create an AI trends report. It should include recent innovations, policy updates, and safety risks. Then synthesize it into a unified summary.\",\n",
    "    expected_output=\"Assign subtasks via the DelegateWorkTool and return a final report.\",\n",
    "    agent=orchestrator,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Crew"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "crew = Crew(\n",
    "    agents=[trend_researcher, policy_analyst, risk_specialist, synthesizer],\n",
    "    tasks=[initial_task],\n",
    "    manager_agent=orchestrator,\n",
    "    verbose=True,\n",
    ")\n",
    "# Run the full workflow\n",
    "result = crew.kickoff()\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fH0uVMgxpLql"
   },
   "source": [
    "### Check your Phoenix project to view the traces and spans from your runs."
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
